import os
import random

from gxl_ai_utils.utils import utils_file

def do_split_scp_for_mk_k2_manifest(input_wav_path, input_text_path, output_manifest_dir,num_split):
    # 切分成100小块
    temp_dir = os.path.join(output_manifest_dir, "temp")
    utils_file.makedir_sil(temp_dir)
    wav_dict = utils_file.load_dict_from_scp(input_wav_path)
    text_dict = utils_file.load_dict_from_scp(input_text_path)
    wav_dict_list = utils_file.do_split_dict(wav_dict, num_split)
    text_dict_list = []
    for wav_dict_i in wav_dict_list:
        text_dict_i = {}
        for key, value in text_dict.items():
            if key in wav_dict_i:
                text_dict_i[key] = value
        text_dict_list.append(text_dict_i)
    for i in utils_file.tqdm(range(num_split), desc="split", total=num_split):
        temp_dir_i = os.path.join(temp_dir, f'temp_{i}')
        utils_file.makedir_sil(temp_dir_i)
        temp_wav_dict = wav_dict_list[i]
        temp_text_dict = text_dict_list[i]
        utils_file.write_dict_to_scp(temp_wav_dict, os.path.join(temp_dir_i, "wav.scp"))
        utils_file.write_dict_to_scp(temp_text_dict, os.path.join(temp_dir_i, "text"))
def do_concat_for_k2_manifest(output_manifest_dir, num_split):
    temp_dir = os.path.join(output_manifest_dir, "temp")
    utils_file.makedir_sil(temp_dir)
    wav_manifest_name, text_manifest_name = utils_file.get_jsonl_filename4icefall("gxldata", "train")
    res_wav_jsonl_list = []
    res_text_jsonl_list = []
    for i in utils_file.tqdm(range(num_split), desc="concat", total=num_split):
        temp_dir_i = os.path.join(temp_dir, f'temp_{i}')
        wav_manifest_path = os.path.join(temp_dir_i, wav_manifest_name)
        text_manifest_path = os.path.join(temp_dir_i, text_manifest_name)
        wav_mani_list = utils_file.load_list_file_clean(wav_manifest_path)
        text_mani_list = utils_file.load_list_file_clean(text_manifest_path)
        res_wav_jsonl_list.extend(wav_mani_list)
        res_text_jsonl_list.extend(text_mani_list)

    output_wav_mani_path = os.path.join(output_manifest_dir, wav_manifest_name)
    output_text_mani_path = os.path.join(output_manifest_dir, text_manifest_name)
    random.shuffle(res_wav_jsonl_list)
    random.shuffle(res_text_jsonl_list)
    utils_file.write_list_to_file(res_wav_jsonl_list, output_wav_mani_path)
    utils_file.write_list_to_file(res_text_jsonl_list, output_text_mani_path)
# 切分成100小块
def do_make_manifest_from_split(output_manifest_dir,num_split):
    temp_dir = os.path.join(output_manifest_dir, "temp")
    utils_file.makedir_sil(temp_dir)
    wav_mani_name, text_mani_name = utils_file.get_jsonl_filename4icefall("gxldata", "train")
    for i in utils_file.tqdm(range(num_split), desc="making", total=num_split):
        temp_dir_i = os.path.join(temp_dir, f'temp_{i}')
        wav_path = os.path.join(temp_dir_i, "wav.scp")
        text_path = os.path.join(temp_dir_i, "text")
        utils_file._do_convert_scp_to_manifest4icefall(
            wav_scp_path=wav_path,
            text_scp_path=text_path,
            wav_manifest_path=os.path.join(temp_dir_i, wav_mani_name),
            text_manifest_path=os.path.join(temp_dir_i, text_mani_name),
            num_thread=8
        )


def main_do_make_manifest(input_wav_path, input_text_path, output_manifest_dir, num_split):
    utils_file.makedir_sil(output_manifest_dir)
    do_split_scp_for_mk_k2_manifest(input_wav_path, input_text_path, output_manifest_dir, num_split)
    do_make_manifest_from_split(output_manifest_dir, num_split)
    do_concat_for_k2_manifest(output_manifest_dir, num_split)


if __name__ == '__main__':
    for part in ["train","dev1","dev2"]:
        input_wav_path = f"/home/work_nfs13/xlgeng/data/seame/{part}/wav.scp"
        input_text_path = f"/home/work_nfs13/xlgeng/data/seame/{part}/text"
        output_manifest_dir = f"/home/work_nfs13/xlgeng/data/seame/{part}/k2_data/train_manifest"
        num_split = 100
        main_do_make_manifest(input_wav_path, input_text_path, output_manifest_dir, num_split)